package com.feidee.fd.sml.algorithm.ml

import com.feidee.fd.sml.algorithm.component.ml.classification.{GBDTComponent, GBDTParam}
import com.feidee.fd.sml.algorithm.forecast.StageFinder
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.apache.spark.ml.classification.GBTClassificationModel
import org.scalatest.FunSuite

/**
  * @Author songhaicheng
  * @Date 2018/08/22
  * @Email: haicheng_song@sui.com
  */
class GBDTComponentSuite extends FunSuite {

  val gbdt = new GBDTComponent
  val paramStr: String =
    """
      |{
      |	'input_pt': '',
      |	'output_pt': '',
      |	'featuresCol': 'features',
      |	'labelCol': 'label',
      | 'probabilityCol': 'pro',
      |	'modelPath': '',
      |	'metrics': ['accuracy', 'auc', 'pr', 'abc']
      |}
    """.stripMargin
  val param: GBDTParam = gbdt.parseParam(new ToolClass().encrypt(paramStr))

  /**
    * GBDT 参数构造方法测试：检测到读取出合法参数 & 成功赋值默认参数，即认为测试通过
    */
  test("GBDT parameter construction test") {
    assert("".equals(param.input_pt) & "".equals(param.output_pt) & "features".equals(param.featuresCol) &
      "label".equals(param.labelCol) & "rawPrediction".equals(param.rawPredictionCol) & "prediction".equals(param.predictionCol) &
      "proto".equals(param.probabilityCol) & param.numIterations == 0 & "".equals(param.modelPath) & param.thresholds.length == 0 &
      param.metrics.sameElements(Array("accuracy", "auc", "pr", "abc"))
    )
  }

  /**
    * GBDT 模型训练功能测试：传入生成的测试数据和正确参数，可以正常训练 GBDT 模型，且预测结果列等于生成数据列数，即认为测试通过
    */
  test("GBDT training function test") {
    // 测试模型训练方法
    val data = TestingDataGenerator.sampleBinaryClassificationData
    val model = gbdt.train(param, data)
    assert(new StageFinder[GBTClassificationModel](model).findWithOrder().get.numClasses == 2)
    assert(model.transform(data).count() == 100)
  }

}
